import pymysql


def checkdb(host, user, pw, dbname):
    """在本地mysql数据库创建初始学生信息"""
    # 打开数据库连接
    db = pymysql.connect(host, user, pw, charset='utf8')
    try:
        # 创建游标，通过连接与数据通信
        cursor = db.cursor()
        # 执行sql语句
        cursor.execute('show databases')
        rows = cursor.fetchall()
        tmp = []
        for row in rows:
            tmp.append("%s" % row)
        # 判断数据库是否存在
        if dbname not in tmp:
            # cursor.execute('drop database if exists ' + dbname)
            cursor.execute('create database if not exists ' + dbname)
            cursor.execute('use ' + dbname)
            #设置所创建的数据库字符集为utf8
            cursor.execute(
                "ALTER DATABASE `%s` CHARACTER SET 'utf8'"
                " COLLATE 'utf8_unicode_ci'" % dbname)
            sql = """
                    create table Student(
                        id int unsigned primary key not null auto_increment,
                        name varchar(100) not null,
                        sex enum('M','F') not null,
                        age tinyint unsigned not null
                    )
                    """
            sql_1 = "insert into student(name,sex,age) " \
                    "values('%s','%s','%d')" % ('张三', 'M', 30)
            sql_2 = "insert into student(name,sex,age) " \
                    "values('%s','%s','%d')" % ('李四', 'F', 20)
            sql_3 = "insert into student(name,sex,age) " \
                    "values('%s','%s','%d')" % ('王五', 'M', 25)
            cursor.execute(sql)
            cursor.execute(sql_1)
            cursor.execute(sql_2)
            cursor.execute(sql_3)
        # 提交到数据库执行
        db.commit()
    except pymysql.Error as e:
        # 事务回滚
        db.rollback()
        print("Mysql Error %d: %s" % (e.args[0], e.args[1]))
    finally:
        # 关闭数据库连接
        db.close()


def deltestDB(host, user, pw, dbname):
    """删除测试数据库"""
    # 打开数据库连接
    db = pymysql.connect(host, user, pw, charset='utf8')
    try:
        # 创建游标，通过连接与数据通信
        cursor = db.cursor()
        # 执行sql语句
        # 删除数据库
        cursor.execute('drop database if exists ' + dbname)
        # 提交到数据库执行
        db.commit()
    except pymysql.Error as e:
        # 事务回滚
        db.rollback()
        print("Mysql Error %d: %s" % (e.args[0], e.args[1]))
    finally:
        # 关闭数据库连接
        db.close()


if __name__ == '__main__':
    checkdb('localhost', 'root', '', 'mystunom')
    key = input('是否要删除测试用数据库？[y/n]:')
    if key == 'y' or key == 'Y':
        deltestDB('localhost', 'root', '', 'mystunom')
